from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import uvicorn
from PyCmpltrtok.common import md5, uuid
import redis
import time
import asyncio
import json

app = FastAPI()
rdb = redis.Redis('127.0.0.1', 6379, 0)
PWD = 'my_password_001'



@app.post("/api")
async def do_infer(request: Request):
    # 接收输入
    req_json = await request.json()  # 请求json
    
    async def get_generator():
        xinput = req_json['input']
        xusername = req_json['username']

        # 防止蹭模型
        xcheck = req_json['check']
        xmy_check = md5(xinput + PWD + xusername)
        if xcheck.lower() != xmy_check.lower():
            print('Checking not passed!')
            return

        print('input:', xinput)

        xoutput = xinput
        for i in range(20):
            n = i + 1
            await asyncio.sleep(0.1)
            xoutput += f'{n}.'
            # 返回输出
            res_dict = dict()
            res_dict['input'] = xinput
            res_dict['output'] = xoutput
            yield json.dumps(res_dict).encode('utf8') + b"\0"

        # 返回输出
        res_dict = dict()
        res_dict['input'] = xinput
        res_dict['output'] = xoutput
        yield json.dumps(res_dict).encode('utf8') + b"\0"
        
    generator = get_generator()
    # https://fastapi.tiangolo.com/advanced/custom-response/#streamingresponse
    return StreamingResponse(generator)


if '__main__' == __name__:
    # https://stackoverflow.com/questions/63177681/is-there-a-difference-between-running-fastapi-from-uvicorn-command-in-dockerfile
    uvicorn.run('fastapi_service:app', host='0.0.0.0', port=7765)
